/*
 * Copyright (C) 2013, Pavel Samko <bulldozerbsg@gmail.com>
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the "Pavel Samko" nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
 * OF SUCH DAMAGE.
 */

#include <windows.h>
#include <wchar.h>
#include <string.h>
#include <assert.h>
#include "main_setup.h"

// WCX export
char *g_plugNamesWcx[] =
{
	// mandatory
	"OpenArchive",
	"ReadHeader",
	"ProcessFile",
	"CloseArchive",
	"SetChangeVolProc",
	"SetProcessDataProc",
	// optional
	"PackFiles",
	"DeleteFiles",
	"GetPackerCaps",
	"ConfigurePacker",
	"StartMemPack",
	"PackToMem",
	"DoneMemPack",
	"CanYouHandleThisFile",
	"PackSetDefaultParams",
	"PkSetCryptCallback",
	"ReadHeaderEx",
	"GetBackgroundFlags",
	// unicode
	"OpenArchiveW",
	"ReadHeaderExW",
	"ProcessFileW",
	"SetChangeVolProcW",
	"SetProcessDataProcW",
	"PackFilesW",
	"DeleteFilesW",
	"StartMemPackW",
	"CanYouHandleThisFileW"
};
size_t g_plugNamesWcxCount = sizeof(g_plugNamesWcx) / sizeof(g_plugNamesWcx[0]);
char g_plugNamesWcxFlag[sizeof(g_plugNamesWcx) / sizeof(g_plugNamesWcx[0])];

WCHAR g_plugPath[MAX_PATH];
WCHAR *g_plugFile;
WCHAR *g_plugExt;

void log(HWND hwnd, WCHAR *s)
{
	int sSize;
	if (s && (sSize = wcslen(s)))
	{
		HWND log = GetDlgItem(hwnd, EDT_LOG);
		int logSize = GetWindowTextLengthW(log);
		int buffCount = sSize + logSize + 1;
		WCHAR *buff = (WCHAR *)malloc(sizeof(WCHAR) * buffCount);
		if (buff)
		{
			GetWindowTextW(log, buff, buffCount);
			wcscat(buff, s);
			SetWindowTextW(log, buff);
			free(buff);
		}
	}
}

int extract_rc_m(HWND hwnd, int rc, BYTE **buff, DWORD *buffSize)
{
	HRSRC res = FindResource(0, MAKEINTRESOURCE(rc), RT_RCDATA);
	HGLOBAL mem = LoadResource(0, res);
	void *data = LockResource(mem);
	DWORD dataSize = SizeofResource(0, res);
	assert(data && dataSize);
	*buff = (BYTE *)malloc(dataSize);
	memcpy(*buff, data, dataSize);
	*buffSize = dataSize;
	return 1;
}

int write_to_file(HWND hwnd, WCHAR *fileName, BYTE *buff, DWORD buffSize)
{
	HANDLE f = CreateFileW(fileName,
		GENERIC_WRITE,
		0,
		NULL,
		CREATE_NEW,
		FILE_ATTRIBUTE_NORMAL,
		NULL);
	while (f == INVALID_HANDLE_VALUE)
	{
		if (GetLastError() != ERROR_FILE_EXISTS)
		{
			return 0;
		}
		static WCHAR msg[2 * MAX_PATH] = {0};
		swprintf(msg, 2 * MAX_PATH, L"File already exists. Overwrite?\n%ls", fileName);
		if (MessageBoxW(hwnd, msg, L"Warning", MB_ICONWARNING | MB_YESNO) != IDYES)
		{
			return 0;
		}
		f = CreateFileW(fileName,
			GENERIC_WRITE,
			0,
			NULL,
			CREATE_ALWAYS,
			FILE_ATTRIBUTE_NORMAL,
			NULL);
	}
	DWORD dataSizeWrited = 0;
	if (!WriteFile(f, buff, buffSize, &dataSizeWrited, NULL) || dataSizeWrited != buffSize)
	{
		CloseHandle(f);
		return 0;
	}
	return 1;
}

int extract_rc_f(HWND hwnd, int rc, WCHAR *fileName)
{
	BYTE *data = NULL;
	DWORD dataSize = 0;
	if (!extract_rc_m(hwnd, rc, &data, &dataSize))
	{
		return 0;
	}
	if (!write_to_file(hwnd, fileName, data, dataSize))
	{
		free(data);
		return 0;
	}
	free(data);
	return 1;
}

size_t addr32_to_offset(DWORD addr, IMAGE_NT_HEADERS32 *peHeader, IMAGE_SECTION_HEADER (*section)[])
{
	for (size_t i = 0; i < peHeader->FileHeader.NumberOfSections; ++i)
	{
		if ((*section)[i].VirtualAddress <= addr && 
			(*section)[i].VirtualAddress + (*section)[i].Misc.VirtualSize >= addr)
		{
			return addr - (*section)[i].VirtualAddress + (*section)[i].PointerToRawData;
		}
	}
	return 0;
}

size_t addr64_to_offset(DWORD addr, IMAGE_NT_HEADERS64 *peHeader, IMAGE_SECTION_HEADER (*section)[])
{
	for (size_t i = 0; i < peHeader->FileHeader.NumberOfSections; ++i)
	{
		if ((*section)[i].VirtualAddress <= addr && 
			(*section)[i].VirtualAddress + (*section)[i].Misc.VirtualSize >= addr)
		{
			return addr - (*section)[i].VirtualAddress + (*section)[i].PointerToRawData;
		}
	}
	return 0;
}

int check_import(HWND hwnd, WCHAR *fileName)
{
	HANDLE f = CreateFileW(fileName,
		GENERIC_READ,
		FILE_SHARE_READ,
		NULL,
		OPEN_EXISTING,
		FILE_ATTRIBUTE_NORMAL,
		NULL);
	if (f == INVALID_HANDLE_VALUE)
	{
		log(hwnd, L"Could not open the file: ");
		log(hwnd, fileName);
		log(hwnd, L"\r\n");
		return 0;
	}
	DWORD fSize = GetFileSize(f, NULL);
	if (fSize == INVALID_FILE_SIZE)
	{
		log(hwnd, L"Invalid file size\r\n");
		CloseHandle(f);
		return 0;
	}
	BYTE *buff = (BYTE *)malloc(fSize);
	if (!buff)
	{
		CloseHandle(f);
		return 0;
	}
	DWORD buffSizeReaded = 0;
	if (!ReadFile(f, buff, fSize, &buffSizeReaded, NULL) || buffSizeReaded != fSize)
	{
		free(buff);
		CloseHandle(f);
		return 0;
	}

	// get and check MZ signature
	IMAGE_DOS_HEADER *mzHeader = (IMAGE_DOS_HEADER *)buff;
	if (mzHeader->e_magic != IMAGE_DOS_SIGNATURE)
	{
		free(buff);
		CloseHandle(f);
		return 0;
	}
	// get PE header and check
	size_t peHeaderOffset = mzHeader->e_lfanew;
	IMAGE_NT_HEADERS32 *peHeader = (IMAGE_NT_HEADERS32 *)(buff + peHeaderOffset);
	if (peHeader->Signature != IMAGE_NT_SIGNATURE ||
		peHeader->FileHeader.Machine != IMAGE_FILE_MACHINE_I386 ||
		peHeader->FileHeader.Characteristics & IMAGE_FILE_DLL == IMAGE_FILE_DLL ||
		!peHeader->FileHeader.SizeOfOptionalHeader ||
		peHeader->FileHeader.SizeOfOptionalHeader > sizeof(IMAGE_OPTIONAL_HEADER32) ||
		peHeader->OptionalHeader.NumberOfRvaAndSizes < IMAGE_DIRECTORY_ENTRY_EXPORT + 1)
	{
		free(buff);
		CloseHandle(f);
		return 0;
	}
	// get section headers
	size_t sectionOffset = peHeaderOffset + sizeof(IMAGE_NT_HEADERS32) - (IMAGE_NUMBEROF_DIRECTORY_ENTRIES - peHeader->OptionalHeader.NumberOfRvaAndSizes) * sizeof(IMAGE_DATA_DIRECTORY);
	//size_t sectionSize = peHeader->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER);
	IMAGE_SECTION_HEADER (*section)[peHeader->FileHeader.NumberOfSections] = (IMAGE_SECTION_HEADER (*)[])(buff + sectionOffset);
	// get export
	size_t exportOffset = addr32_to_offset(peHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress, peHeader, section);
	if (!exportOffset)
	{
		free(buff);
		CloseHandle(f);
		return 0;
	}
	//size_t exportSize = peHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].Size;
	IMAGE_EXPORT_DIRECTORY *export = (IMAGE_EXPORT_DIRECTORY *)(buff + exportOffset);
	// get name table
	size_t namesOffset = addr32_to_offset(export->AddressOfNames, peHeader, section);
	DWORD (*names)[export->NumberOfNames] = (DWORD (*)[])(buff + namesOffset);

	// WCX
	memset(g_plugNamesWcxFlag, 0, sizeof(g_plugNamesWcxFlag));

	char *name = NULL;
	for (size_t i = 0; i < export->NumberOfNames; ++i)
	{
		name = (char *)(buff + addr32_to_offset((*names)[i], peHeader, section));
		if (!name)
		{
			continue;
		}
		// WCX
		for (size_t j = 0; j < g_plugNamesWcxCount; ++j)
		{
			if (!strcmp(name, g_plugNamesWcx[j]))
			{
				g_plugNamesWcxFlag[j] = 1;
			}
		}
	}

	free(buff);
	CloseHandle(f);
	return 1;
}

int fix_export(HWND hwnd, WCHAR *fileName, BYTE *buff, DWORD buffSize)
{
	// get and check MZ signature
	IMAGE_DOS_HEADER *mzHeader = (IMAGE_DOS_HEADER *)buff;
	if (mzHeader->e_magic != IMAGE_DOS_SIGNATURE)
	{
		free(buff);
		return 0;
	}
	// get PE header and check
	size_t peHeaderOffset = mzHeader->e_lfanew;
	IMAGE_NT_HEADERS64 *peHeader = (IMAGE_NT_HEADERS64 *)(buff + peHeaderOffset);
	if (peHeader->Signature != IMAGE_NT_SIGNATURE ||
		peHeader->FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64 ||
		peHeader->FileHeader.Characteristics & IMAGE_FILE_DLL == IMAGE_FILE_DLL ||
		!peHeader->FileHeader.SizeOfOptionalHeader ||
		peHeader->FileHeader.SizeOfOptionalHeader > sizeof(IMAGE_OPTIONAL_HEADER64) ||
		peHeader->OptionalHeader.NumberOfRvaAndSizes < IMAGE_DIRECTORY_ENTRY_EXPORT + 1)
	{
		free(buff);
		return 0;
	}
	// get section headers
	size_t sectionOffset = peHeaderOffset + sizeof(IMAGE_NT_HEADERS64) - (IMAGE_NUMBEROF_DIRECTORY_ENTRIES - peHeader->OptionalHeader.NumberOfRvaAndSizes) * sizeof(IMAGE_DATA_DIRECTORY);
	//size_t sectionSize = peHeader->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER);
	IMAGE_SECTION_HEADER (*section)[peHeader->FileHeader.NumberOfSections] = (IMAGE_SECTION_HEADER (*)[])(buff + sectionOffset);
	// get export
	size_t exportOffset = addr64_to_offset(peHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress, peHeader, section);
	if (!exportOffset)
	{
		free(buff);
		return 0;
	}
	//size_t exportSize = peHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].Size;
	IMAGE_EXPORT_DIRECTORY *export = (IMAGE_EXPORT_DIRECTORY *)(buff + exportOffset);
	// get name table
	size_t namesOffset = addr64_to_offset(export->AddressOfNames, peHeader, section);
	size_t ordinalsOffset = addr64_to_offset(export->AddressOfNameOrdinals, peHeader, section);
	DWORD (*names)[export->NumberOfNames] = (DWORD (*)[])(buff + namesOffset);
	WORD (*ordinals)[export->NumberOfNames] = (WORD (*)[])(buff + ordinalsOffset);

	char *name = NULL;
	for (size_t i = 0; i < export->NumberOfNames;)
	{
		name = (char *)(buff + addr64_to_offset((*names)[i], peHeader, section));
		if (!name)
		{
			++i;
			continue;
		}
		// WCX
		int removed = 0;
		for (size_t j = 0; j < g_plugNamesWcxCount; ++j)
		{
			if (!strcmp(name, g_plugNamesWcx[j]) && !g_plugNamesWcxFlag[j])
			{
				removed = 1;
				--export->NumberOfNames;
				for(size_t k = i; k < export->NumberOfNames; ++k)
				{
					(*names)[k] = (*names)[k + 1];
					(*ordinals)[k] = (*ordinals)[k + 1];
				}
				break;
			}
		}
		if (!removed)
		{
			++i;
			continue;
		}
	}

	if (!write_to_file(hwnd, fileName, buff, buffSize))
	{
		free(buff);
		return 0;
	}

	free(buff);
	return 1;
}

BOOL CALLBACK dlg_proc(HWND hwnd, UINT Message, WPARAM wParam, LPARAM lParam)
{
	BOOL result = FALSE;
	switch (Message)
	{
		case WM_INITDIALOG:
			result = TRUE;
			break;
		case WM_COMMAND:
			switch (LOWORD(wParam))
			{
			case BTN_SELECT:
				{
					g_plugPath[0] = 0;
					OPENFILENAMEW ofn;
					ZeroMemory(&ofn, sizeof(OPENFILENAMEW));
					ofn.lStructSize = sizeof(OPENFILENAMEW);
					ofn.hwndOwner = hwnd;
					ofn.lpstrFilter = L"All types of plugins (*.wcx; *.wfx; *.wlx; *.wdx)\0*.wcx;*.wfx;*.wlx;*.wdx\0";
					ofn.lpstrFile = g_plugPath;
					ofn.nMaxFile = sizeof(g_plugPath);
					ofn.lpstrTitle = L"Select Total Commander plugin";
					ofn.Flags = OFN_HIDEREADONLY | OFN_FILEMUSTEXIST;

					if (GetOpenFileNameW(&ofn))
					{
						log(hwnd, L"Selected file: ");
						log(hwnd, g_plugPath);
						log(hwnd, L"\r\nCheck file...\r\n");
						
						if (check_import(hwnd, g_plugPath))
						{
							EnableWindow(GetDlgItem(hwnd, BTN_SETUP), TRUE);
							g_plugFile = wcsrchr(g_plugPath, L'\\');
							if (g_plugFile)
							{
								g_plugFile[0] = 0;
								++g_plugFile;
							}
							g_plugExt = wcsrchr(g_plugFile, L'.');
							{
								g_plugExt[0] = 0;
								++g_plugExt;
							}
						} else
						{
							log(hwnd, L"\r\n");
						}
					}
				}
				break;
			case BTN_SETUP:
				{
					// extract bridge files from res
					int setup = 1;
					static WCHAR fileName[MAX_PATH] = {0};

					swprintf(fileName, MAX_PATH, L"%ls\\%ls.exe", g_plugPath, g_plugFile);
					log(hwnd, L"Extract file: ");
					log(hwnd, fileName);
					log(hwnd, L"... ");
					setup &= extract_rc_f(hwnd, 1, fileName);
					log(hwnd, setup? L"Success.\r\n" : L"Failed!\r\n");

					swprintf(fileName, MAX_PATH, L"%ls\\%ls.%ls64", g_plugPath, g_plugFile, g_plugExt);
					log(hwnd, L"Extract file: ");
					log(hwnd, fileName);
					log(hwnd, L"... ");
					BYTE *data = NULL;
					DWORD dataSize = 0;
					setup &= extract_rc_m(hwnd, 2, &data, &dataSize);
					log(hwnd, setup? L"Success.\r\n" : L"Failed!\r\n");

					if (setup && fix_export(hwnd, fileName, data, dataSize))
					{
						log(hwnd, L"The program has been successfully installed.\r\n\r\n");
					} else
					{
						log(hwnd, L"Process failed.\r\n\r\n");
					}
				}
				break;
			case IDCANCEL:
				EndDialog(hwnd, IDCANCEL);
				break;	
			}
			result = TRUE;
			break;
		default:
			result = FALSE;
	}
	return result;
}

int CALLBACK WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow)
{
	DialogBoxW(GetModuleHandle(NULL), MAKEINTRESOURCEW(DLG_BRIDGE), NULL, dlg_proc);
	return 0;
}
